import mindspore


def custom_Cifar10Dataset_generator(transform, batch_size):
    training_data = mindspore.dataset.Cifar10Dataset(dataset_dir='data', usage='train').map(operations=transform)
    testing_data = mindspore.dataset.Cifar10Dataset(dataset_dir='data', usage='test').map(operations=transform)

    train_data = training_data.batch(batch_size=batch_size).shuffle(buffer_size=4).batch(drop_remainder=True)
    test_data = testing_data.batch(batch_size=batch_size).shuffle(buffer_size=4).batch(drop_remainder=True)
    return train_data, test_data


transform = mindspore.dataset.transforms.Compose([mindspore.dataset.vision.ToTensor(),
                                                  mindspore.dataset.vision.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
                                                  mindspore.dataset.vision.Resize((224, 224))
                                                  ])

batch_size = 64
train_data, test_data = custom_Cifar10Dataset_generator(transform, batch_size)